import graphscope
from graphscope.dataset import load_p2p_network

from graphscope.analytical.udf.decorators import pregel
from graphscope.framework.app import AppAssets

import numpy as np

# 装饰器, 定义顶点数据和消息数据的类型
# debug 看过程 还要测试无向图能不能用
@pregel(vd_type="double", md_type="double")
class SSSP_Pregel(AppAssets):
    @staticmethod
    def Init(v, context):
        v.set_value(1000000000.0)

    @staticmethod
    def Compute(messages, v, context):
        src_id = context.get_config(b"src")
        cur_dist = v.value()
        new_dist = 1000000000.0
        if v.id() == src_id:
            new_dist = 0
        for message in messages:
            new_dist = min(message, new_dist)
        if new_dist < cur_dist:
            v.set_value(new_dist)
            for e_label_id in range(context.edge_label_num()):
                edges = v.outgoing_edges(e_label_id)
                for e in edges:
                    v.send(e.vertex(), new_dist + e.get_int(2))
        v.vote_to_halt()

    @staticmethod
    def Combine(messages):
        ret = 1000000000.0
        for m in messages:
            ret = min(ret, m)
        return ret


def create_acline(R: float, X: float, Vmax: float, Vmin: float, Imax: float) -> dict:
    return {"R": R, "X": X, "V2max": Vmax**2, "V2min": Vmin**2, "I2max": Imax**2}


@pregel(vd_type="double", md_type="double")
class SSSP_Pregel2(AppAssets):
    @staticmethod
    def Init(v, context):
        v.set_value(1000000000.0)

    @staticmethod
    def Compute(messages, v, context):
        src_id = context.get_config(b"src")
        t_args = context.get_config(b"terminal")
        cur_dist = v.value()
        new_dist = 1000000000.0
        if v.id() == src_id:
            new_dist = 0
        for message in messages:
            new_dist = min(message, new_dist)
        if new_dist < cur_dist:
            v.set_value(new_dist)
            for e_label_id in range(context.edge_label_num()):
                edges = v.outgoing_edges(e_label_id)
                for e in edges:
                    v.send(e.vertex(), new_dist + e.get_int(2))
        v.vote_to_halt()

    @staticmethod
    def Combine(messages):
        ret = 1000000000.0
        for m in messages:
            ret = min(ret, m)
        return ret


g = load_p2p_network(generate_eid=False)

# 加载自己的算法
my_app = SSSP_Pregel()

# 在图上运行自己的算法，得到计算结果
# 这里 `src` 是与s`context.get_config(b"src")` 相对应的
ctx = my_app(g, src="6")
r2 = (
    ctx.to_dataframe({"node": "v:host.id", "r": "r:host"})
    .sort_values(by=["node"])
    .to_numpy(dtype=float)
)
my_app2 = SSSP_Pregel2()
# ret = my_app2(g, src="6", terminal=terminal_args(70))
ret = my_app2(g, src="6", terminal=[1,10,20])
ret = my_app2(g, src="6", terminal={"L":[2,3],"P":[30 ,40]})
ret = my_app2(g, src="6", terminal=create_acline)
